#!/usr/bin/env python

# Dependencies:
# Windows version of du from Sysinternals (NOT THE CYGWIN VERSION)
# Cairo (available from GTK) http://www.gtk.org/download-windows.html
# PyCairo
# CairoPlot https://launchpad.net/cairoplot
# PIL http://www.pythonware.com/products/pil/

# April 2010
# Author: George Louthan http://georgerloutha.nthefourth.com
# License: BSD-new

from CairoPlot import *
import os
from PIL import Image,ImageFont,ImageDraw
from datetime import datetime
from ConfigParser import ConfigParser
#from ctypes import ctypes

DISK_DRIVE = 'C' # hex
#DISK_DRIVE = 'E' # fatbot

AFFILIATION_TEXT = 'Univ. of Tulsa iSec'
DEVICE_TEXT = 'Fatbot'
IN_FILE = 'du.exe'
ARCHIVE_PATH='archive'

#TOTAL_SPACE = 500103639 # fatbot
TOTAL_SPACE = 319243153 # hex
IMAGE_PATH='.'
SHARE_IMAGE_NAME = 'shares.png'
SHARE_FREE_IMAGE_NAME = 'shares_free.png'
TIME_STRING = datetime.now().isoformat().split('.')[0]

WIDTH=500
HEIGHT=300

def main():    
    
    config = ConfigParser()
    config.read('dugraph.conf')
    global DISK_DRIVE, TOTAL_SPACE, IN_FILE, ARCHIVE_PATH, IMAGE_PATH, \
            AFFILIATION_TEXT, DEVICE_TEXT, WIDTH, HEIGHT
    
    DISK_DRIVE = config.get('disk','drive')
    TOTAL_SPACE = config.getint('disk','totspace')
    IN_FILE = config.get('dufile','file')
    ARCHIVE_PATH = config.get('dufile','archive')
    IMAGE_PATH = config.get('output','path')
    AFFILIATION_TEXT = config.get('output','affiliation')
    DEVICE_TEXT = config.get('output','device')
    WIDTH=config.getint('output','width')
    HEIGHT=config.getint('output','height')
    
    used_space = 0
    images_generated = []  
    
    diskuse = open('du.dat')
    lines = []
    for l in diskuse:
        lines+=[l]
    
    directories = {}
    shares = {}
    share_subs = {}
    
    for l in lines:
        if l == '\r\n' or l.split()[0]=='Files:':
            break
        path_parts = ' '.join(l.split()[1:]).split('\\')
        if path_parts[1]=='':
            continue
        if len(path_parts) == 2: # top level, meaning "share-level"
            shares[path_parts[1]] = int(''.join(l.split()[0].split(','))) #strip commas
        else: # below top level, meaning sub-share-level
            if path_parts[1] not in share_subs:
                share_subs[path_parts[1]] = {path_parts[2] : int(''.join(l.split()[0].split(',')))}
            else:
                share_subs[path_parts[1]][path_parts[2]] = int(''.join(l.split()[0].split(',')))
    
    for s in shares:
        used_space += shares[s]
    
    images_generated+=[SHARE_IMAGE_NAME]
    print 'Generating'
    donut_plot(os.path.join(IMAGE_PATH,SHARE_IMAGE_NAME),shares,WIDTH,HEIGHT, gradient=True, background=(.100,.100,.200))
    mark_image(os.path.join(IMAGE_PATH,SHARE_IMAGE_NAME),'Fileshare overview')
    for s in share_subs:
        print 'Generating',os.path.join(IMAGE_PATH,s+'.png')
        donut_plot(os.path.join(IMAGE_PATH,s+'.png'),share_subs[s],WIDTH,HEIGHT, gradient=True, background=(.100,.100,.200))
        mark_image(os.path.join(IMAGE_PATH,s+'.png'),'Share: ' + s)
    shares['empty'] = TOTAL_SPACE - used_space
    images_generated += [SHARE_FREE_IMAGE_NAME]
    print 'Generating', SHARE_FREE_IMAGE_NAME
    donut_plot(os.path.join(IMAGE_PATH,SHARE_FREE_IMAGE_NAME),shares,WIDTH,HEIGHT, gradient=True, background=(.100,.100,.200))
    mark_image(os.path.join(IMAGE_PATH,SHARE_FREE_IMAGE_NAME),'Free space overview')

def mark_image(img,detail_text):
    text = AFFILIATION_TEXT + '\nDevice: ' + DEVICE_TEXT + '\n' + detail_text
    text_list = text.split('\n')
    y=0
    i = Image.open(img)
    dr = ImageDraw.Draw(i)
    for line in text_list:
        dr.text((5,y),line)
        y+=10
    dr.text((WIDTH-48,HEIGHT-20),TIME_STRING.split('T')[1])
    dr.text((WIDTH-60,HEIGHT-10),TIME_STRING.split('T')[0])
    i.save(img,'PNG')

if __name__ == '__main__':
    main()